# -*- coding: utf-8 -*-
"""
Created on Mon Dec 11 11:27:32 2017

@author: za-xuzhaoye
"""
import tushare as ts
import pandas as pd
import numpy as np
from pandas import Series
import math 
#from arr import code

#basicinfo=ts.get_stock_basics()
#idlist=Series(basicinfo.index)
#
#for aa in range(0,1):
#    SecurityId= idlist[aa]
#    mktdata = ts.get_k_data(SecurityId,start='2017-01-01',end='2018-01-08').sort_index()
#    popen = mktdata['open'].values
#    pclose = mktdata['close'].values
#    phigh = mktdata['high'].values
#    plow = mktdata['low'].values
#   

def RSRS(SecurityId,mktdata,popen,pclose,phigh,plow):
    '''calculate RSRS signal'''
    window = 18
    a = np.array([np.nan] * len(mktdata))
    b = [np.nan] * len(mktdata)  # If betas required.
    c = [np.nan] * len(mktdata) 
    high = mktdata.high.values
    low = mktdata[['low']].assign(constant=1).values
    
    for i in range(window, len(mktdata)):
        y = high[(i - window):i]
        X = low[(i - window):i]
        # betas = Inverse(X'.X).X'.y
        betas = np.linalg.inv(X.T.dot(X)).dot(X.T).dot(y)
        y_hat = betas.dot(low[i, :])
        a[i] = y_hat
        b[i] = betas.tolist()[0]  # real betas 
        c[i]= betas.tolist()[1] #constant term 
        
        b= Series(b)
        mean=pd.rolling_mean(b,window=window)
        std=pd.rolling_std(b,window=window)
        up=mean+std
        down=mean-std
        
    return b,up,down
    
def signal_construction(mktdata,SecurityId,b,up,down,pclose,popen):    
    '''strategy: if RSRS signal is higer than up, long the equity, if RSRS is lower than down, short the equity   '''     
        
    tmp1=b>up
    tmp2=b<down
    
    sgn_long=np.where(tmp1==True,1,0)
    sgn_short=np.where(tmp2==True,-1,0)
    sgn_hd = np.zeros(len(sgn_long))
    index_long=np.zeros(len(sgn_long))
    index_short=np.zeros(len(sgn_long))
    holdprice=np.zeros(len(sgn_long)) #当前持有的成本
    rt=np.zeros(len(sgn_long))
    pnl=np.zeros(len(sgn_long))#profit and loss
    countlong=np.zeros(len(sgn_long))
    longsgnrate=np.zeros(len(sgn_long))
    countshort=np.zeros(len(sgn_long))
    shortsgnrate=np.zeros(len(sgn_long))
 
    '''construct tdbook, which contains oder's direction and price'''
    for ii in np.arange(len(mktdata)):
        if ii>0:
                if (sgn_long[ii-1] == 1):
                    sgn_hd[ii] = sgn_long[ii-1] 
                if (sgn_hd[ii-1]==1)and (sgn_short[ii-1]==0):
                    sgn_hd[ii]=sgn_hd[ii-1]    
                if  (sgn_short[ii-1]==-1):
                    sgn_hd[ii] = sgn_short[ii-1]     
                if(sgn_hd[ii-1]==-1)and (sgn_long[ii-1]==0):
                    sgn_hd[ii]=sgn_hd[ii-1]    
                    
                if(sgn_hd[ii-1]<sgn_hd[ii]):
                  index_long[ii] = - np.sign(sgn_hd[ii]) * popen[ii] 
                else:
                 index_long[ii]=0
            
                if(sgn_hd[ii-1]>sgn_hd[ii]):
                  index_short[ii] = - np.sign(sgn_hd[ii]) * pclose[ii] 
                else:
                 index_short[ii]=0
      
        #define holdprice          
                if (sgn_hd[ii]==1) and (index_long[ii]<0):
                    holdprice[ii]=index_long[ii]
                if (sgn_hd[ii]==1) and (index_long[ii]==0):
                    holdprice[ii]=holdprice[ii-1]
                    
                if (sgn_hd[ii]==-1)and(index_short[ii]>0):
                    holdprice[ii]=index_short[ii]
                if (sgn_hd[ii]==-1)and(index_short[ii]==0):
                    holdprice[ii]=holdprice[ii-1] 
        # calculate pnl
                if (sgn_hd[ii]>sgn_hd[ii-1]):
                    pnl[ii]=0
                    if (holdprice[ii]==holdprice[ii-1]):
                        pnl[ii]=0
                        
                if (sgn_hd[ii]<sgn_hd[ii-1]):
                    pnl[ii]=holdprice[ii]+holdprice[ii-1]
                
                rt[ii]=pclose[ii]/pclose[ii-1]-1
                
                TestTerm=5
                sgn_long1=sgn_long.tolist()
                TestRange=sgn_long1[(ii-TestTerm):(ii)]
                countlong[ii] =np.count_nonzero(TestRange)
                longsgnrate[ii]=countlong[ii]/TestTerm
                
                sgn_short1=sgn_short.tolist()
                TestShortRange=sgn_short1[(ii-TestTerm):(ii)]
                countshort[ii] =np.count_nonzero(TestShortRange)
                shortsgnrate[ii]=countshort[ii]/TestTerm
                
                        
    d = {'sgn_long':sgn_long,'sgn_short':sgn_short,'sgn_hd':sgn_hd,'holdprice':holdprice,
         'pnl':pnl,'pclose':pclose,'rt':rt,'longsgnrate':longsgnrate,'shortsgnrate':shortsgnrate}
    tdbook=pd.DataFrame(d, columns=['sgn_long','sgn_short','sgn_hd','holdprice','pnl',
                                    'pclose','rt','longsgnrate','shortsgnrate'])
    tdbook.insert(0,column='datetime',value=mktdata['date'])
#    tdbook.insert(0,column='name',value=basicinfo['name'])
#    tdbook.insert(0,column='industry',value=basicinfo['industry'])
    
    return tdbook, sgn_hd, holdprice,longsgnrate,shortsgnrate

#b,up,down=RSRS(SecurityId,mktdata,popen,pclose,phigh,plow)
#tdbook, sgn_hd, holdprice,longsgnrate,shortsgnrate=signal_construction(basicinfo,mktdata,SecurityId,b,up,down,pclose,popen)


#%%
def CalculateResult(mktdata,SecurityId,tdbook,sgn_hd,holdprice,pclose,capital=1000000,unit=100,
                    rate=0.0025, year=250, rf=0.0035):

    OrderedDict={'StartTime':np.nan,'EndTime':np.nan,'Profit':np.nan,'AnnualReturn':np.nan,
                 'AnnualVol':np.nan,'SharpeRatio':np.nan,'Max_Drawdown':np.nan,}
    
    if len(tdbook)==0:
        OrderedDict['ID']= np.nan
        OrderedDict['name'] = np.String
        OrderedDict['industry'] = np.String
        OrderedDict['Profit']=np.nan
        OrderedDict['AnnualReturn'] = np.nan
        OrderedDict['AnnualVol'] = np.nan
        OrderedDict['SharpeRatio'] = np.nan
        OrderedDict['Max_Drawdown'] = np.nan
        OrderedDict['StartTime'] = np.nan
        OrderedDict['EndTime'] = np.nan
        OrderedDict['ProfitRate'] = np.nan
        OrderedDict['Position'] = np.nan
        
    else:
        
        size = np.zeros(len(sgn_hd))                             
        Cap = np.zeros(len(sgn_hd)) + capital  
        cost= np.zeros(len(sgn_hd))
        RT=np.zeros(len(sgn_hd))
        
        sgn_hd=tdbook['sgn_hd']
        
        for jj in np.arange(len(sgn_hd)):
          if jj>0: 
                if (sgn_hd[jj]>sgn_hd[jj-1]):
                     size[jj]=np.floor(Cap[jj-1]/abs(holdprice[jj])/unit/(1+rate))
                     cost[jj]=abs(holdprice[jj])*size[jj]*unit*rate
                     Cap[jj]=Cap[jj-1]-size[jj]*unit*(1+rate)*abs(holdprice[jj])
                     RT[jj]= (Cap[jj]+abs(holdprice[jj])*size[jj]*unit)/Cap[jj-1]-1
                 
                if (sgn_hd[jj]==1)and(sgn_hd[jj]==sgn_hd[jj-1]):
                     size[jj]=size[jj-1]
                     cost[jj]=0
                     Cap[jj]=Cap[jj-1]
                     RT[jj]= pclose[jj]/pclose[jj-1]-1
                 
                if (sgn_hd[jj]<sgn_hd[jj-1]):
                     size[jj]=0
                     cost[jj]=abs(holdprice[jj])*size[jj-1]*unit*rate
                     Cap[jj]=Cap[jj-1]+cost[jj]/rate-cost[jj]
                     RT[jj]=Cap[jj]/(Cap[jj-1]+pclose[jj-1]*size[jj-1]*unit)-1
                 
                if (sgn_hd[jj]== -1) and (sgn_hd[jj]==sgn_hd[jj-1]):
                     size[jj]=size[jj-1]
                     cost[jj]=0
                     Cap[jj]=Cap[jj-1]
                     RT[jj]=0
               
        tdbook1={'sgn_hd':sgn_hd,'Cap':Cap,'size':size,'cost':cost,'return':RT}        
        TradeRecord = pd.DataFrame(tdbook1,columns=['sgn_hd','Cap','size','cost','return'])
        TradeRecord.insert(0,column='datetime',value=mktdata['date'])  
    
       
        Profit=(TradeRecord['Cap'].iloc[-1])+(TradeRecord['size'].iloc[-1])*(tdbook['pclose'].iloc[-1])*unit-(TradeRecord['Cap'].iloc[0])
        ProfitRate=Profit/(TradeRecord['Cap'].iloc[0])
        StartTime=tdbook['datetime'].iloc[0]
        EndTime=tdbook['datetime'].iloc[-1]
        Max_Drawdown= TradeRecord['return'].min()
        AnnualReturn=(Profit/(TradeRecord['Cap'].iloc[0])*year)/len(mktdata)
        AnnualVol=TradeRecord['return'].std()*math.sqrt(year)
        SharpeRatio=(AnnualReturn-rf)/AnnualVol
        Holding_position=tdbook['sgn_hd'].iloc[-1]
        LongTermHoldingReturn=pclose[-1]/pclose[0]-1
        Longsgnrate=tdbook['longsgnrate'].iloc[-1]
        shortsgnrate=tdbook['shortsgnrate'].iloc[-1]
#        name=basicinfo['pe']
#        industry=basicinfo['industry']
        
        OrderedDict['ID']= SecurityId
#        OrderedDict['name'] = name
#        OrderedDict['industry'] =industry
        OrderedDict['Profit'] = Profit
        OrderedDict['StartTime'] = StartTime
        OrderedDict['EndTime'] = EndTime
        OrderedDict['ProfitRate'] = ProfitRate
        OrderedDict['AnnualReturn'] = AnnualReturn
        OrderedDict['AnnualVol'] = AnnualVol
        OrderedDict['SharpeRatio'] = SharpeRatio
        OrderedDict['Max_Drawdown'] = Max_Drawdown
        OrderedDict['Position'] =  Holding_position
        OrderedDict['LongTermHoldingReturn']=LongTermHoldingReturn
        OrderedDict['Longsgnrate']=Longsgnrate
        OrderedDict['shortsgnrate']=shortsgnrate
        
        Result= pd.DataFrame(OrderedDict,columns=['ID','Profit','StartTime','EndTime','ProfitRate','AnnualReturn','AnnualVol',
                                                  'SharpeRatio','Max_Drawdown','Position','LongTermHoldingReturn','Longsgnrate','shortsgnrate'],index=[0])
        
    return TradeRecord,Result
#
#b,up,down=RSRS(SecurityId,mktdata,popen,pclose,phigh,plow)
#tdbook, sgn_hd, holdprice,longsgnrate,shortsgnrate=signal_construction(mktdata,SecurityId,b,up,down,pclose,popen)
#TradeRecord,Result=CalculateResult(basicinfo,mktdata,SecurityId,tdbook,sgn_hd,holdprice,pclose,capital=1000000,unit=100,rate=0.0025, year=270, rf=0.0035)
